"""
重试机制 - 提供多种退避策略的自动重试功能
支持指数退避、线性退避、固定延迟等策略
"""

import asyncio
import time
import random
import functools
from typing import (
    Callable, Optional, Union, List, Type, Any, Dict,
    Awaitable, TypeVar, Protocol, runtime_checkable
)
from dataclasses import dataclass, field
from enum import Enum
from abc import ABC, abstractmethod
from datetime import datetime, timedelta

from ..logging.manager import get_logger
from ..events.bus import get_event_bus, Event

# 类型变量
F = TypeVar('F', bound=Callable[..., Any])
AsyncF = TypeVar('AsyncF', bound=Callable[..., Awaitable[Any]])

class RetryStrategy(Enum):
    """重试策略"""
    EXPONENTIAL = "exponential"   # 指数退避
    LINEAR = "linear"             # 线性退避
    FIXED = "fixed"               # 固定延迟
    CUSTOM = "custom"             # 自定义策略

@runtime_checkable
class BackoffStrategy(Protocol):
    """退避策略接口"""
    
    def get_delay(self, attempt: int) -> float:
        """获取重试延迟时间"""
        ...

class ExponentialBackoff:
    """指数退避策略"""
    
    def __init__(
        self,
        initial_delay: float = 1.0,
        max_delay: float = 60.0,
        multiplier: float = 2.0,
        jitter: bool = True
    ):
        self.initial_delay = initial_delay
        self.max_delay = max_delay
        self.multiplier = multiplier
        self.jitter = jitter
    
    def get_delay(self, attempt: int) -> float:
        """计算指数退避延迟"""
        delay = self.initial_delay * (self.multiplier ** (attempt - 1))
        delay = min(delay, self.max_delay)
        
        if self.jitter:
            # 添加随机抖动，避免惊群效应
            delay = delay * (0.5 + random.random() * 0.5)
        
        return delay

class LinearBackoff:
    """线性退避策略"""
    
    def __init__(
        self,
        initial_delay: float = 1.0,
        increment: float = 1.0,
        max_delay: float = 60.0,
        jitter: bool = True
    ):
        self.initial_delay = initial_delay
        self.increment = increment
        self.max_delay = max_delay
        self.jitter = jitter
    
    def get_delay(self, attempt: int) -> float:
        """计算线性退避延迟"""
        delay = self.initial_delay + (self.increment * (attempt - 1))
        delay = min(delay, self.max_delay)
        
        if self.jitter:
            delay = delay * (0.8 + random.random() * 0.4)
        
        return delay

class FixedBackoff:
    """固定延迟策略"""
    
    def __init__(self, delay: float = 1.0, jitter: bool = True):
        self.delay = delay
        self.jitter = jitter
    
    def get_delay(self, attempt: int) -> float:
        """计算固定延迟"""
        delay = self.delay
        
        if self.jitter:
            delay = delay * (0.8 + random.random() * 0.4)
        
        return delay

@dataclass
class RetryConfig:
    """重试配置"""
    max_attempts: int = 3
    backoff_strategy: BackoffStrategy = field(default_factory=ExponentialBackoff)
    retryable_exceptions: List[Type[Exception]] = field(default_factory=lambda: [Exception])
    stop_on_exceptions: List[Type[Exception]] = field(default_factory=list)
    retry_on_result: Optional[Callable[[Any], bool]] = None
    timeout: Optional[float] = None
    
    def should_retry(self, exception: Exception, result: Any = None) -> bool:
        """判断是否应该重试"""
        # 检查是否是停止重试的异常
        for stop_exception in self.stop_on_exceptions:
            if isinstance(exception, stop_exception):
                return False
        
        # 检查是否是可重试的异常
        if exception:
            for retryable_exception in self.retryable_exceptions:
                if isinstance(exception, retryable_exception):
                    return True
            return False
        
        # 检查结果是否需要重试
        if self.retry_on_result and result is not None:
            return self.retry_on_result(result)
        
        return False

class RetryManager:
    """重试管理器"""
    
    def __init__(self):
        self.logger = get_logger("retry_manager")
        self.event_bus = get_event_bus()
        self.retry_configs: Dict[str, RetryConfig] = {}
    
    def register_config(self, name: str, config: RetryConfig):
        """注册重试配置"""
        self.retry_configs[name] = config
        self.logger.info(f"Registered retry config: {name}")
    
    def get_config(self, name: str) -> Optional[RetryConfig]:
        """获取重试配置"""
        return self.retry_configs.get(name)
    
    async def retry_async(
        self,
        func: Callable[..., Awaitable[Any]],
        *args,
        config: RetryConfig = None,
        config_name: str = None,
        **kwargs
    ) -> Any:
        """异步函数重试"""
        # 获取配置
        if config_name and config_name in self.retry_configs:
            config = self.retry_configs[config_name]
        elif config is None:
            config = RetryConfig()
        
        attempt = 0
        last_exception = None
        start_time = time.time()
        
        while attempt < config.max_attempts:
            attempt += 1
            
            # 检查超时
            if config.timeout and (time.time() - start_time) > config.timeout:
                self.logger.warning(f"Retry timeout after {time.time() - start_time:.2f}s")
                break
            
            try:
                self.logger.debug(f"Attempting function call (attempt {attempt}/{config.max_attempts})")
                
                # 执行函数
                if config.timeout:
                    remaining_time = config.timeout - (time.time() - start_time)
                    result = await asyncio.wait_for(func(*args, **kwargs), timeout=remaining_time)
                else:
                    result = await func(*args, **kwargs)
                
                # 发布重试成功事件
                if attempt > 1:
                    self.event_bus.publish(Event(
                        event_type="retry.success",
                        data={
                            "function": f"{func.__module__}.{func.__name__}",
                            "attempt": attempt,
                            "total_attempts": config.max_attempts,
                            "duration": time.time() - start_time
                        }
                    ))
                
                return result
                
            except Exception as e:
                last_exception = e
                
                # 判断是否应该重试
                if not config.should_retry(e):
                    self.logger.warning(f"Non-retryable exception: {type(e).__name__}: {e}")
                    break
                
                if attempt < config.max_attempts:
                    # 计算延迟时间
                    delay = config.backoff_strategy.get_delay(attempt)
                    
                    self.logger.warning(
                        f"Function call failed (attempt {attempt}/{config.max_attempts}), "
                        f"retrying in {delay:.2f}s: {type(e).__name__}: {e}"
                    )
                    
                    # 发布重试事件
                    self.event_bus.publish(Event(
                        event_type="retry.attempt",
                        data={
                            "function": f"{func.__module__}.{func.__name__}",
                            "attempt": attempt,
                            "total_attempts": config.max_attempts,
                            "delay": delay,
                            "exception": str(e),
                            "exception_type": type(e).__name__
                        }
                    ))
                    
                    # 等待重试
                    await asyncio.sleep(delay)
                else:
                    self.logger.error(
                        f"Function call failed after {attempt} attempts: {type(e).__name__}: {e}"
                    )
        
        # 发布重试失败事件
        self.event_bus.publish(Event(
            event_type="retry.failed",
            data={
                "function": f"{func.__module__}.{func.__name__}",
                "total_attempts": attempt,
                "max_attempts": config.max_attempts,
                "duration": time.time() - start_time,
                "last_exception": str(last_exception) if last_exception else None
            }
        ))
        
        # 抛出最后一个异常
        if last_exception:
            raise last_exception
        else:
            raise RuntimeError("All retry attempts failed")
    
    def retry_sync(
        self,
        func: Callable[..., Any],
        *args,
        config: RetryConfig = None,
        config_name: str = None,
        **kwargs
    ) -> Any:
        """同步函数重试"""
        # 获取配置
        if config_name and config_name in self.retry_configs:
            config = self.retry_configs[config_name]
        elif config is None:
            config = RetryConfig()
        
        attempt = 0
        last_exception = None
        start_time = time.time()
        
        while attempt < config.max_attempts:
            attempt += 1
            
            # 检查超时
            if config.timeout and (time.time() - start_time) > config.timeout:
                self.logger.warning(f"Retry timeout after {time.time() - start_time:.2f}s")
                break
            
            try:
                self.logger.debug(f"Attempting function call (attempt {attempt}/{config.max_attempts})")
                
                # 执行函数
                result = func(*args, **kwargs)
                
                # 发布重试成功事件
                if attempt > 1:
                    self.event_bus.publish(Event(
                        event_type="retry.success",
                        data={
                            "function": f"{func.__module__}.{func.__name__}",
                            "attempt": attempt,
                            "total_attempts": config.max_attempts,
                            "duration": time.time() - start_time
                        }
                    ))
                
                return result
                
            except Exception as e:
                last_exception = e
                
                # 判断是否应该重试
                if not config.should_retry(e):
                    self.logger.warning(f"Non-retryable exception: {type(e).__name__}: {e}")
                    break
                
                if attempt < config.max_attempts:
                    # 计算延迟时间
                    delay = config.backoff_strategy.get_delay(attempt)
                    
                    self.logger.warning(
                        f"Function call failed (attempt {attempt}/{config.max_attempts}), "
                        f"retrying in {delay:.2f}s: {type(e).__name__}: {e}"
                    )
                    
                    # 发布重试事件
                    self.event_bus.publish(Event(
                        event_type="retry.attempt",
                        data={
                            "function": f"{func.__module__}.{func.__name__}",
                            "attempt": attempt,
                            "total_attempts": config.max_attempts,
                            "delay": delay,
                            "exception": str(e),
                            "exception_type": type(e).__name__
                        }
                    ))
                    
                    # 等待重试
                    time.sleep(delay)
                else:
                    self.logger.error(
                        f"Function call failed after {attempt} attempts: {type(e).__name__}: {e}"
                    )
        
        # 发布重试失败事件
        self.event_bus.publish(Event(
            event_type="retry.failed",
            data={
                "function": f"{func.__module__}.{func.__name__}",
                "total_attempts": attempt,
                "max_attempts": config.max_attempts,
                "duration": time.time() - start_time,
                "last_exception": str(last_exception) if last_exception else None
            }
        ))
        
        # 抛出最后一个异常
        if last_exception:
            raise last_exception
        else:
            raise RuntimeError("All retry attempts failed")

# 全局重试管理器
_retry_manager: Optional[RetryManager] = None

def get_retry_manager() -> RetryManager:
    """获取全局重试管理器"""
    global _retry_manager
    if _retry_manager is None:
        _retry_manager = RetryManager()
    return _retry_manager

def set_retry_manager(manager: RetryManager):
    """设置全局重试管理器"""
    global _retry_manager
    _retry_manager = manager

# 装饰器
def retry(
    max_attempts: int = 3,
    backoff_strategy: BackoffStrategy = None,
    retryable_exceptions: List[Type[Exception]] = None,
    stop_on_exceptions: List[Type[Exception]] = None,
    retry_on_result: Optional[Callable[[Any], bool]] = None,
    timeout: Optional[float] = None,
    config_name: str = None
):
    """重试装饰器（同步函数）"""
    def decorator(func: F) -> F:
        # 创建重试配置
        config = RetryConfig(
            max_attempts=max_attempts,
            backoff_strategy=backoff_strategy or ExponentialBackoff(),
            retryable_exceptions=retryable_exceptions or [Exception],
            stop_on_exceptions=stop_on_exceptions or [],
            retry_on_result=retry_on_result,
            timeout=timeout
        )
        
        # 注册配置（如果提供了名称）
        if config_name:
            get_retry_manager().register_config(config_name, config)
        
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return get_retry_manager().retry_sync(
                func, *args, config=config, config_name=config_name, **kwargs
            )
        
        return wrapper
    
    return decorator

def retry_async(
    max_attempts: int = 3,
    backoff_strategy: BackoffStrategy = None,
    retryable_exceptions: List[Type[Exception]] = None,
    stop_on_exceptions: List[Type[Exception]] = None,
    retry_on_result: Optional[Callable[[Any], bool]] = None,
    timeout: Optional[float] = None,
    config_name: str = None
):
    """重试装饰器（异步函数）"""
    def decorator(func: AsyncF) -> AsyncF:
        # 创建重试配置
        config = RetryConfig(
            max_attempts=max_attempts,
            backoff_strategy=backoff_strategy or ExponentialBackoff(),
            retryable_exceptions=retryable_exceptions or [Exception],
            stop_on_exceptions=stop_on_exceptions or [],
            retry_on_result=retry_on_result,
            timeout=timeout
        )
        
        # 注册配置（如果提供了名称）
        if config_name:
            get_retry_manager().register_config(config_name, config)
        
        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            return await get_retry_manager().retry_async(
                func, *args, config=config, config_name=config_name, **kwargs
            )
        
        return wrapper
    
    return decorator

# 预定义配置
QUICK_RETRY = RetryConfig(
    max_attempts=3,
    backoff_strategy=FixedBackoff(delay=0.5),
    timeout=5.0
)

STANDARD_RETRY = RetryConfig(
    max_attempts=5,
    backoff_strategy=ExponentialBackoff(initial_delay=1.0, max_delay=30.0),
    timeout=60.0
)

PERSISTENT_RETRY = RetryConfig(
    max_attempts=10,
    backoff_strategy=ExponentialBackoff(initial_delay=2.0, max_delay=300.0),
    timeout=600.0
)

# 示例使用
if __name__ == "__main__":
    import asyncio
    import random
    
    # 注册预定义配置
    manager = get_retry_manager()
    manager.register_config("quick", QUICK_RETRY)
    manager.register_config("standard", STANDARD_RETRY)
    manager.register_config("persistent", PERSISTENT_RETRY)
    
    # 使用装饰器
    @retry(max_attempts=3, backoff_strategy=ExponentialBackoff())
    def unreliable_function():
        if random.random() < 0.7:  # 70%概率失败
            raise ConnectionError("Network connection failed")
        return "Success!"
    
    @retry_async(max_attempts=5, config_name="standard")
    async def unreliable_async_function():
        await asyncio.sleep(0.1)
        if random.random() < 0.6:  # 60%概率失败
            raise TimeoutError("Request timeout")
        return "Async Success!"
    
    # 手动重试
    async def manual_retry_example():
        try:
            result = await manager.retry_async(
                unreliable_async_function,
                config_name="quick"
            )
            print(f"Result: {result}")
        except Exception as e:
            print(f"Failed after all retries: {e}")
    
    # 运行示例
    async def main():
        # 测试同步重试
        try:
            result = unreliable_function()
            print(f"Sync result: {result}")
        except Exception as e:
            print(f"Sync failed: {e}")
        
        # 测试异步重试
        await manual_retry_example()
    
    asyncio.run(main())
